/**********
This library is free software; you can redistribute it and/or modify it under
the terms of the GNU Lesser General Public License as published by the
Free Software Foundation; either version 3 of the License, or (at your
option) any later version. (See <http://www.gnu.org/copyleft/lesser.html>.)

This library is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for
more details.

You should have received a copy of the GNU Lesser General Public License
along with this library; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA
**********/
// "liveMedia"
// Copyright (c) 1996-2018 Live Networks, Inc.  All rights reserved.
// A class that encapsulates MPEG-2 Transport Stream 'index files'/
// These index files are used to implement 'trick play' operations
// (seek-by-time, fast forward, reverse play) on Transport Stream files.
//
// Implementation

#include "include/MPEG2TransportStreamIndexFile.hh"
#include "include/InputFile.hh"

MPEG2TransportStreamIndexFile
::MPEG2TransportStreamIndexFile(UsageEnvironment &env, char const *indexFileName)
        : Medium(env),
          fFileName(strDup(indexFileName)), fFid(NULL), fMPEGVersion(0), fCurrentIndexRecordNum(0),
          fCachedPCR(0.0f), fCachedTSPacketNumber(0), fNumIndexRecords(0) {
    // Get the file size, to determine how many index records it contains:
    u_int64_t indexFileSize = GetFileSize(indexFileName, NULL);
    if (indexFileSize % INDEX_RECORD_SIZE != 0) {
        env << "Warning: Size of the index file \"" << indexFileName
            << "\" (" << (unsigned) indexFileSize
            << ") is not a multiple of the index record size ("
            << INDEX_RECORD_SIZE << ")\n";
    }
    fNumIndexRecords = (unsigned long) (indexFileSize / INDEX_RECORD_SIZE);
}

MPEG2TransportStreamIndexFile *MPEG2TransportStreamIndexFile
::createNew(UsageEnvironment &env, char const *indexFileName) {
    if (indexFileName == NULL) return NULL;
    MPEG2TransportStreamIndexFile *indexFile
            = new MPEG2TransportStreamIndexFile(env, indexFileName);

    // Reject empty or non-existent index files:
    if (indexFile->getPlayingDuration() == 0.0f) {
        delete indexFile;
        indexFile = NULL;
    }

    return indexFile;
}

MPEG2TransportStreamIndexFile::~MPEG2TransportStreamIndexFile() {
    closeFid();
    delete[] fFileName;
}

void MPEG2TransportStreamIndexFile
::lookupTSPacketNumFromNPT(float &npt, unsigned long &tsPacketNumber,
                           unsigned long &indexRecordNumber) {
    if (npt <= 0.0 || fNumIndexRecords == 0) { // Fast-track a common case:
        npt = 0.0f;
        tsPacketNumber = indexRecordNumber = 0;
        return;
    }

    // If "npt" is the same as the one that we last looked up, return its cached result:
    if (npt == fCachedPCR) {
        tsPacketNumber = fCachedTSPacketNumber;
        indexRecordNumber = fCachedIndexRecordNumber;
        return;
    }

    // Search for the pair of neighboring index records whose PCR values span "npt".
    // Use the 'regula-falsi' method.
    Boolean success = False;
    unsigned long ixFound = 0;
    do {
        unsigned long ixLeft = 0, ixRight = fNumIndexRecords - 1;
        float pcrLeft = 0.0f, pcrRight;
        if (!readIndexRecord(ixRight)) break;
        pcrRight = pcrFromBuf();
        if (npt > pcrRight) npt = pcrRight;
        // handle "npt" too large by seeking to the last frame of the file

        while (ixRight - ixLeft > 1 && pcrLeft < npt && npt <= pcrRight) {
            unsigned long ixNew = ixLeft
                                  + (unsigned long) (((npt - pcrLeft) / (pcrRight - pcrLeft)) *
                                                     (ixRight - ixLeft));
            if (ixNew == ixLeft || ixNew == ixRight) {
                // use bisection instead:
                ixNew = (ixLeft + ixRight) / 2;
            }
            if (!readIndexRecord(ixNew)) break;
            float pcrNew = pcrFromBuf();
            if (pcrNew < npt) {
                pcrLeft = pcrNew;
                ixLeft = ixNew;
            } else {
                pcrRight = pcrNew;
                ixRight = ixNew;
            }
        }
        if (ixRight - ixLeft > 1 || npt <= pcrLeft || npt > pcrRight)
            break; // bad PCR values in index file?

        ixFound = ixRight;
        // "Rewind' until we reach the start of a Video Sequence or GOP header:
        success = rewindToCleanPoint(ixFound);
    } while (0);

    if (success && readIndexRecord(ixFound)) {
        // Return (and cache) information from record "ixFound":
        npt = fCachedPCR = pcrFromBuf();
        tsPacketNumber = fCachedTSPacketNumber = tsPacketNumFromBuf();
        indexRecordNumber = fCachedIndexRecordNumber = ixFound;
    } else {
        // An error occurred: Return the default values, for npt == 0:
        npt = 0.0f;
        tsPacketNumber = indexRecordNumber = 0;
    }
    closeFid();
}

void MPEG2TransportStreamIndexFile
::lookupPCRFromTSPacketNum(unsigned long &tsPacketNumber, Boolean reverseToPreviousCleanPoint,
                           float &pcr, unsigned long &indexRecordNumber) {
    if (tsPacketNumber == 0 || fNumIndexRecords == 0) { // Fast-track a common case:
        pcr = 0.0f;
        indexRecordNumber = 0;
        return;
    }

    // If "tsPacketNumber" is the same as the one that we last looked up, return its cached result:
    if (tsPacketNumber == fCachedTSPacketNumber) {
        pcr = fCachedPCR;
        indexRecordNumber = fCachedIndexRecordNumber;
        return;
    }

    // Search for the pair of neighboring index records whose TS packet #s span "tsPacketNumber".
    // Use the 'regula-falsi' method.
    Boolean success = False;
    unsigned long ixFound = 0;
    do {
        unsigned long ixLeft = 0, ixRight = fNumIndexRecords - 1;
        unsigned long tsLeft = 0, tsRight;
        if (!readIndexRecord(ixRight)) break;
        tsRight = tsPacketNumFromBuf();
        if (tsPacketNumber > tsRight) tsPacketNumber = tsRight;
        // handle "tsPacketNumber" too large by seeking to the last frame of the file

        while (ixRight - ixLeft > 1 && tsLeft < tsPacketNumber && tsPacketNumber <= tsRight) {
            unsigned long ixNew = ixLeft
                                  + (unsigned long) (
                    ((tsPacketNumber - tsLeft) / (tsRight - tsLeft)) * (ixRight - ixLeft));
            if (ixNew == ixLeft || ixNew == ixRight) {
                // Use bisection instead:
                ixNew = (ixLeft + ixRight) / 2;
            }
            if (!readIndexRecord(ixNew)) break;
            unsigned long tsNew = tsPacketNumFromBuf();
            if (tsNew < tsPacketNumber) {
                tsLeft = tsNew;
                ixLeft = ixNew;
            } else {
                tsRight = tsNew;
                ixRight = ixNew;
            }
        }
        if (ixRight - ixLeft > 1 || tsPacketNumber <= tsLeft || tsPacketNumber > tsRight)
            break; // bad PCR values in index file?

        ixFound = ixRight;
        if (reverseToPreviousCleanPoint) {
            // "Rewind' until we reach the start of a Video Sequence or GOP header:
            success = rewindToCleanPoint(ixFound);
        } else {
            success = True;
        }
    } while (0);

    if (success && readIndexRecord(ixFound)) {
        // Return (and cache) information from record "ixFound":
        pcr = fCachedPCR = pcrFromBuf();
        fCachedTSPacketNumber = tsPacketNumFromBuf();
        if (reverseToPreviousCleanPoint) tsPacketNumber = fCachedTSPacketNumber;
        indexRecordNumber = fCachedIndexRecordNumber = ixFound;
    } else {
        // An error occurred: Return the default values, for tsPacketNumber == 0:
        pcr = 0.0f;
        indexRecordNumber = 0;
    }
    closeFid();
}

Boolean MPEG2TransportStreamIndexFile
::readIndexRecordValues(unsigned long indexRecordNum,
                        unsigned long &transportPacketNum, u_int8_t &offset,
                        u_int8_t &size, float &pcr, u_int8_t &recordType) {
    if (!readIndexRecord(indexRecordNum)) return False;

    transportPacketNum = tsPacketNumFromBuf();
    offset = offsetFromBuf();
    size = sizeFromBuf();
    pcr = pcrFromBuf();
    recordType = recordTypeFromBuf();
    return True;
}

float MPEG2TransportStreamIndexFile::getPlayingDuration() {
    if (fNumIndexRecords == 0 || !readOneIndexRecord(fNumIndexRecords - 1)) return 0.0f;

    return pcrFromBuf();
}

int MPEG2TransportStreamIndexFile::mpegVersion() {
    if (fMPEGVersion != 0) return fMPEGVersion; // we already know it

    // Read the first index record, and figure out the MPEG version from its type:
    if (!readOneIndexRecord(0)) return 0; // unknown; perhaps the indecx file is empty?

    setMPEGVersionFromRecordType(recordTypeFromBuf());
    return fMPEGVersion;
}

Boolean MPEG2TransportStreamIndexFile::openFid() {
    if (fFid == NULL && fFileName != NULL) {
        if ((fFid = OpenInputFile(envir(), fFileName)) != NULL) {
            fCurrentIndexRecordNum = 0;
        }
    }

    return fFid != NULL;
}

Boolean MPEG2TransportStreamIndexFile::seekToIndexRecord(unsigned long indexRecordNumber) {
    if (!openFid()) return False;

    if (indexRecordNumber == fCurrentIndexRecordNum) return True; // we're already there

    if (SeekFile64(fFid, (int64_t) (indexRecordNumber * INDEX_RECORD_SIZE), SEEK_SET) != 0)
        return False;
    fCurrentIndexRecordNum = indexRecordNumber;
    return True;
}

Boolean MPEG2TransportStreamIndexFile::readIndexRecord(unsigned long indexRecordNum) {
    do {
        if (!seekToIndexRecord(indexRecordNum)) break;
        if (fread(fBuf, INDEX_RECORD_SIZE, 1, fFid) != 1) break;
        ++fCurrentIndexRecordNum;

        return True;
    } while (0);

    return False; // an error occurred
}

Boolean MPEG2TransportStreamIndexFile::readOneIndexRecord(unsigned long indexRecordNum) {
    Boolean result = readIndexRecord(indexRecordNum);
    closeFid();

    return result;
}

void MPEG2TransportStreamIndexFile::closeFid() {
    if (fFid != NULL) {
        CloseInputFile(fFid);
        fFid = NULL;
    }
}

float MPEG2TransportStreamIndexFile::pcrFromBuf() {
    unsigned pcr_int = (fBuf[5] << 16) | (fBuf[4] << 8) | fBuf[3];
    u_int8_t pcr_frac = fBuf[6];
    return pcr_int + pcr_frac / 256.0f;
}

unsigned long MPEG2TransportStreamIndexFile::tsPacketNumFromBuf() {
    return (fBuf[10] << 24) | (fBuf[9] << 16) | (fBuf[8] << 8) | fBuf[7];
}

void MPEG2TransportStreamIndexFile::setMPEGVersionFromRecordType(u_int8_t recordType) {
    if (fMPEGVersion != 0) return; // we already know it

    u_int8_t const recordTypeWithoutStartBit = recordType & ~0x80;
    if (recordTypeWithoutStartBit >= 1 && recordTypeWithoutStartBit <= 4) fMPEGVersion = 2;
    else if (recordTypeWithoutStartBit >= 5 && recordTypeWithoutStartBit <= 10) fMPEGVersion = 5;
        // represents H.264
    else if (recordTypeWithoutStartBit >= 11 && recordTypeWithoutStartBit <= 16) fMPEGVersion = 6;
    // represents H.265
}

Boolean MPEG2TransportStreamIndexFile::rewindToCleanPoint(unsigned long &ixFound) {
    Boolean success = False; // until we learn otherwise

    while (ixFound > 0) {
        if (!readIndexRecord(ixFound)) break;

        u_int8_t recordType = recordTypeFromBuf();
        setMPEGVersionFromRecordType(recordType);

        // A 'clean point' is the start of a 'frame' from which a decoder can cleanly resume
        // handling the stream.  For H.264, this is a SPS.  For H.265, this is a VPS.
        // For MPEG-2, this is a Video Sequence Header, or a GOP.

        if ((recordType & 0x80) != 0) { // This is the start of a 'frame'
            recordType &= ~0x80; // remove the 'start of frame' bit
            if (fMPEGVersion == 5) { // H.264
                if (recordType == 5/*SPS*/) {
                    success = True;
                    break;
                }
            } else if (fMPEGVersion == 6) { // H.265
                if (recordType == 11/*VPS*/) {
                    success = True;
                    break;
                }
            } else { // MPEG-1, 2, or 4
                if (recordType == 1/*VSH*/) {
                    success = True;
                    break;
                } else if (recordType == 2/*GOP*/) {
                    // Hack: If the preceding record is for a Video Sequence Header, then use it instead:
                    unsigned long newIxFound = ixFound;

                    while (--newIxFound > 0) {
                        if (!readIndexRecord(newIxFound)) break;
                        recordType = recordTypeFromBuf();
                        if ((recordType & 0x7F) != 1) break; // not a Video Sequence Header
                        if ((recordType & 0x80) != 0) { // this is the start of the VSH; use it
                            ixFound = newIxFound;
                            break;
                        }
                    }
                }
                success = True;
                break;
            }
        }

        // Keep checking, from the previous record:
        --ixFound;
    }
    if (ixFound == 0) success = True; // use record 0 anyway

    return success;
}
